#include "recon.h"

recon::recon()
{
    String dirPath = QCoreApplication::applicationDirPath().toStdString() + "/images";
    map<string, Mat> images = findImagesInDiretcory(dirPath);
    extractFeatures(images);
    matchFeatures();
    buildTracks();
    reconstructFromTracks();
    visualize3D();
}

map<string, Mat> recon::findImagesInDiretcory(string dirPath)
{
    std::vector<string> src_name;
    map<string, Mat> images1;
    cv::glob(dirPath, src_name, false);
    if(src_name.empty())
    {
        std::cout<<"load images error! The size is zero!"<<std::endl;
    }
    else
    {
        std::cout<<"load images success! The size is "+ std::to_string(src_name.size())<<std::endl;
        std::sort(src_name.begin(), src_name.end());
        for (const auto& i : src_name) {
            imagesNames.push_back(i);
            //map<string, Mat> images;
            images[i] = imread(i);
            images1[i] = imread(i);
            imageIDs[i] = images.size() - 1;
        }
    }
    return images1;
}

void recon::extractFeatures(map<string, Mat> images)
{
    auto detector = AKAZE::create();
    auto extractor = AKAZE::create();

    for (const auto& i : images) {
        vector<KeyPoint> kp;
        Mat dp;
        Mat grayscale;

        cvtColor(i.second, grayscale, COLOR_BGR2GRAY);
        detector->detect(grayscale, kp);
        extractor->compute(grayscale, kp, dp);

        keypoints[i.first] = kp;
        descriptors[i.first] = dp;
    }
}

vector<DMatch> recon::matchWithRatioTest(const DescriptorMatcher& matcher, Mat desc1, Mat desc2, float MATCH_RATIO_THRESHOLD)
{
    vector<vector<DMatch>> nnMatch;
    matcher.knnMatch(desc1, desc2, nnMatch, 2);

    // Ratio test filter
    vector<DMatch> ratioMatched;
    for (size_t i = 0; i < nnMatch.size(); i++) {
        DMatch first = nnMatch[i][0];
        float dist1 = nnMatch[i][0].distance;
        float dist2 = nnMatch[i][1].distance;
        if (dist1 < MATCH_RATIO_THRESHOLD * dist2) {
            ratioMatched.push_back(first);
        }
    }
    return ratioMatched;
}

void recon::matchFeatures()
{
    map<pair<string, string>, vector<DMatch>> matchsRes;

    size_t imagesNums = keypoints.size();
    BFMatcher matcher(NORM_HAMMING);
    for (size_t i = 0; i < imagesNums - 1; i++) {
        for (size_t j = i + 1; j < imagesNums; j++) {
            // Match with ratio test filter
            string iKey = imagesNames[i];
            string jKey = imagesNames[j];
            vector<DMatch> match
                = matchWithRatioTest(matcher, descriptors[iKey], descriptors[jKey]);
            // Reciprocity test filter
            vector<DMatch> matchRcp
                = matchWithRatioTest(matcher, descriptors[jKey], descriptors[iKey]);
            vector<DMatch> merged;
            for (const DMatch& dmr : matchRcp) {
                bool found = false;
                for (const DMatch& dm : match) {
                    // Only accept match if 1 matches 2 AND 2 matches 1.
                    if (dmr.queryIdx == dm.trainIdx and dmr.trainIdx == dm.queryIdx) {
                        merged.push_back(dm);
                        found = true;
                        break;
                    }
                }
                if (found) {
                    continue;
                }
            }

            matches[make_pair(iKey, jKey)] = merged;
            std::cout<<merged.size()<<"  ";
        }
        std::cout<<std::endl;
    }
}

void recon::buildTracks()
{
    using namespace boost;

    struct ImageFeature {
        string image;
        size_t featureID;
    };
    typedef adjacency_list<listS, vecS, undirectedS, ImageFeature> Graph;
    typedef graph_traits<Graph>::vertex_descriptor Vertex;

    map<pair<string, int>, Vertex> vertexByImageFeature;

    Graph g;

    //map<string, vector<KeyPoint>> keypoints;
    //map<string, Mat> descriptors;
    // Add vertices - image features
    for (const auto& imgi : keypoints) {
        for (size_t i = 0; i < imgi.second.size(); i++) {
            Vertex v = add_vertex(g);
            g[v].image = imgi.first;
            g[v].featureID = i;
            vertexByImageFeature[make_pair(imgi.first, i)] = v;
        }
    }

    //map<pair<string, string>, vector<DMatch>> matches;
    // Add edges - feature matches
    for (const auto& match : matches) {
        for (const DMatch& dm : match.second) {
            Vertex& vI = vertexByImageFeature[make_pair(match.first.first, dm.queryIdx)];
            Vertex& vJ = vertexByImageFeature[make_pair(match.first.second, dm.trainIdx)];
            add_edge(vI, vJ, g);
        }
    }

    using Filtered = filtered_graph<Graph, keep_all, std::function<bool(Vertex)>>;
    Filtered gFiltered(g, keep_all {}, [&g](Vertex vd) { return degree(vd, g) > 0; });

    // Get connected components
    std::vector<int> component(num_vertices(gFiltered), -1);
    int num = connected_components(gFiltered, &component[0]);
    map<int, vector<Vertex>> components;
    for (size_t i = 0; i != component.size(); ++i) {
        if (component[i] >= 0) {
            components[component[i]].push_back(i);
        }
    }
    // Filter bad components (with more than 1 feature from a single image)
    std::vector<int> vertexInGoodComponent(num_vertices(gFiltered), -1);
    map<int, vector<Vertex>> goodComponents;
    for (const auto& c : components) {
        set<string> imagesInComponent;
        bool isComponentGood = true;
        for (int j = 0; j < c.second.size(); ++j) {
            const string imgId = g[c.second[j]].image;
            if (imagesInComponent.count(imgId) > 0) {
                // Image already represented in this component
                isComponentGood = false;
                break;
            } else {
                imagesInComponent.insert(imgId);
            }
        }
        if (isComponentGood) {
            for (int j = 0; j < c.second.size(); ++j) {
                vertexInGoodComponent[c.second[j]] = 1;
            }
            goodComponents[c.first] = c.second;
        }
    }

    Filtered gGoodComponents(g, keep_all {},
        [&vertexInGoodComponent](Vertex vd) { return vertexInGoodComponent[vd] > 0; });


    const int accum = std::accumulate(goodComponents.begin(), goodComponents.end(), 0,
        [](int a, pair<const int, vector<Vertex>>& v) { return a + v.second.size(); });

    if (saveDebugVisualizations) {
        struct my_node_writer {
            my_node_writer(Graph& g_, const map<string, int>& iid_)
                : g(g_)
                , iid(iid_) {};
            void operator()(std::ostream& out, Vertex v)
            {
                const int imgId = iid[g[v].image];
                out << " [label=\"" << imgId
                    << "\" colorscheme=\"accent8\" fillcolor=" << (imgId + 1)
                    << " style=filled]";
            };
            Graph g;
            map<string, int> iid;
        };
        std::ofstream ofs("match_graph_good_components.dot");
        write_graphviz(ofs, gGoodComponents, my_node_writer(g, imageIDs));
        std::ofstream ofsf("match_graph_filtered.dot");
        write_graphviz(ofsf, gFiltered, my_node_writer(g, imageIDs));
    }

    // Each component is a track
    const size_t nViews = imagesNames.size();
    tracks.resize(nViews);
    for (int i = 0; i < nViews; i++) {
        tracks[i].create(2, goodComponents.size(), CV_64FC1);
        tracks[i].setTo(-1.0);
    }
    int i = 0;
    for (auto c = goodComponents.begin(); c != goodComponents.end(); ++c, ++i) {
        for (const int v : c->second) {
            const int imageID = imageIDs[g[v].image];
            const size_t featureID = g[v].featureID;
            const Point2f p = keypoints[g[v].image][featureID].pt;
            tracks[imageID].at<double>(0, i) = p.x;
            tracks[imageID].at<double>(1, i) = p.y;
        }
    }

    if (saveDebugVisualizations) {
        vector<Scalar> colors
            = { CV_RGB(240, 248, 255), CV_RGB(250, 235, 215), CV_RGB(0, 255, 255),
                  CV_RGB(127, 255, 212), CV_RGB(240, 255, 255), CV_RGB(245, 245, 220),
                  CV_RGB(255, 228, 196), CV_RGB(255, 235, 205), CV_RGB(0, 0, 255),
                  CV_RGB(138, 43, 226), CV_RGB(165, 42, 42), CV_RGB(222, 184, 135) };

        vector<Mat> imagesM;
        for (const auto m : images)
            imagesM.push_back(m.second);
        Mat out;
        hconcat(vector<Mat>(imagesM.begin(), imagesM.begin() + 4), out);
        RNG& rng = cv::theRNG();
        const Size imgS = imagesM[0].size();
        for (int tId = 0; tId < 20; tId++) {
            const int trackId = rng(tracks[0].cols); // Randomize a track ID

            // Show track over images
            for (int i = 0; i < 3; i++) {
                Point2f a = Point2f(tracks[i].col(trackId));
                Point2f b = Point2f(tracks[i + 1].col(trackId));

                if (a.x < 0 or a.y < 0 or b.x < 0 or b.y < 0) {
                    continue;
                }

                const Scalar c = colors[tId % colors.size()];
                a.x += imgS.width * i;
                b.x += imgS.width * (i + 1);
                circle(out, a, 7, c, FILLED);
                circle(out, b, 7, c, FILLED);
                line(out, a, b, c, 3);
            }
            imwrite("tracks.jpg", out);

            // Show track patches
            const int patchSize = 20;
            const Point2f patch(patchSize, patchSize);
            for (int i = 0; i < tracks.size(); i++) {
                Point2f a = Point2f(tracks[i].col(trackId));
                if (a.x < patchSize or a.y < patchSize or a.x > imgS.width - patchSize
                    or a.y > imgS.height - patchSize) {
                    continue;
                }

                imwrite("track_" + to_string(trackId) + "_" + to_string(i) + ".png",
                    imagesM[i](Rect(a - patch, a + patch)));
            }
        }
    }
}

bool recon::reconstructFromTracks()
{
    const Size imgS = images.begin()->second.size();
    const float f = std::max(imgS.width, imgS.height);
    Mat K
        = Mat(Matx33f { f, 0.0, imgS.width / 2.0f, 0.0, f, imgS.height / 2.0f, 0.0, 0.0, 1.0 });

    cv::sfm::reconstruct(tracks, Rs, Ts, K, points3d, true);

    K.copyTo(K_);

    if (Rs.size() != imagesNames.size()) {
        return false;
    }

    if (tracks[0].cols != points3d.size()) {

    }

    // Create the point cloud
    pointCloud.clear();
    for (const auto& p : points3d)
        pointCloud.emplace_back(Vec3f(p));

    // Get the point colors
    pointCloudColor.resize(pointCloud.size(), Vec3b(0, 255, 0));
    vector<Point2f> point2d(1);
    for (int i = 0; i < (int)pointCloud.size(); i++) {
        for (int j = 0; j < imagesNames.size(); ++j) {
            Mat point3d = Mat(pointCloud[i]).reshape(1, 1);
            cv::projectPoints(point3d, Rs[j], Ts[j], K_, Mat(), point2d);
            if (point2d[0].x < 0 or point2d[0].x >= imgS.width or point2d[0].y < 0
                or point2d[0].y >= imgS.height) {
                continue;
            }
            pointCloudColor[i] = images[imagesNames[j]].at<Vec3b>(point2d[0]);
            break;
        }
    }
    return true;
}

void recon::visualize3D()
{
    ofstream dataFile;
    string path = QCoreApplication::applicationDirPath().toStdString() + "/dataFile.txt";
    dataFile.open(path, ofstream::app);
    // 朝TXT文档中写入数据
    std::cout<<pointCloud.size()<<std::endl;
    for(auto & p : pointCloud){
        dataFile<<p(0)<<" "<<p(1)<<" "<<p(2)<<endl;
    }
    // 关闭文档
    dataFile.close();
    std::cout<<"save ok"<<std::endl;
}








